from numpy import array
from OpenGL.GL import *

class Axes(object):
    """
    Axes makes it easy to render coordinate axes etc. with OpenGL.
    """

    def __init__(self):
        super(Axes, self).__init__()
        # store some data to use later
        self.vshader_src = """
            #version 110
            uniform mat4 pvTransform;
            void main()
            {
                gl_Position = pvTransform * gl_Vertex;
                gl_FrontColor = gl_Color;
                gl_BackColor = gl_Color;
            }
            """
        self.fshader_src = """
            #version 110
            void main()
            {
                gl_FragColor = gl_Color;
            }
            """
        self.vertex = array([0, 0, 0, 1,
                             1, 0, 0, 1,
                             0, 0, 0, 1,
                             0, 1, 0, 1,
                             0, 0, 0, 1,
                             0, 0, 1, 1,
                             -1, 0, -1, 1,
                              1, 0, -1, 1,
                              1, 0,  1, 1,
                             -1, 0,  1, 1], dtype='float32')
        self.color = array([1, 0, 0, 1,
                            1, 0, 0, 1,
                            0, 1, 0, 1,
                            0, 1, 0, 1,
                            0, 0, 1, 1,
                            0, 0, 1, 1,
                            0.2, 0.2, 0.2, 1,
                            0.2, 0.2, 0.2, 1,
                            0.2, 0.2, 0.2, 1,
                            0.2, 0.2, 0.2, 1], dtype='float32')
        self.edge = array([0, 1,
                           2, 3,
                           4, 5,
                           6, 7,
                           7, 8,
                           8, 9,
                           9, 6], dtype='uint32')
        # save GL-initialization status
        self.glStatus = False

    def initialized(self):
        return self.glStatus

    def initializeGL(self):
        # first compile shader program
        vshader_id = glCreateShader(GL_VERTEX_SHADER)
        glShaderSource(vshader_id, self.vshader_src)
        glCompileShader(vshader_id)

        fshader_id = glCreateShader(GL_FRAGMENT_SHADER)
        glShaderSource(fshader_id, self.fshader_src)
        glCompileShader(fshader_id)

        self.shader = glCreateProgram()
        glAttachShader(self.shader, vshader_id)
        glAttachShader(self.shader, fshader_id)
        glLinkProgram(self.shader)

        self.pvLocation = glGetUniformLocation(self.shader, "pvTransform")

        # then set up VBO stuff
        self.meshbuffers = glGenBuffers(3)
        glBindBuffer(GL_ARRAY_BUFFER, self.meshbuffers[0])
        glBufferData(GL_ARRAY_BUFFER, self.vertex, GL_STATIC_DRAW)
        glBindBuffer(GL_ARRAY_BUFFER, self.meshbuffers[1])
        glBufferData(GL_ARRAY_BUFFER, self.color, GL_STATIC_DRAW)
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, self.meshbuffers[2])
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, self.edge, GL_STATIC_DRAW)

        # remember we've done this
        self.glStatus = True

    def boundingBox(self):
        return [-1, 0, -1],[1, 1, 1]

    def paintGL(self, pvTransform):
        glUseProgram(self.shader)
        glUniformMatrix4fv(self.pvLocation, 1, True, pvTransform.flatten())
        glEnableClientState(GL_VERTEX_ARRAY)
        glBindBuffer(GL_ARRAY_BUFFER, self.meshbuffers[0])
        glVertexPointer(4, GL_FLOAT, 0, None)
        glEnableClientState(GL_COLOR_ARRAY)
        glBindBuffer(GL_ARRAY_BUFFER, self.meshbuffers[1])
        glColorPointer(4, GL_FLOAT, 0, None)
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, self.meshbuffers[2])
        glDrawElements(GL_LINES, len(self.edge), GL_UNSIGNED_INT, None)
        glDisableClientState(GL_COLOR_ARRAY)
        glDisableClientState(GL_VERTEX_ARRAY)

